package com.hyj.langchain4j_springboot.config;

import com.alibaba.dashscope.threads.runs.Run;
import com.hyj.langchain4j_springboot.service.ToolsService;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.*;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class AiConfig {

    public interface Assistant{
        String chat(String message);
        //流式响应
        TokenStream stream(String message);
        //预设角色流式响应
        @SystemMessage("""
                你是东方航空公司的客户聊天支持代理，请以友好的的方式来进行回复。
                你正在通过在线聊天系统与客户互动。
                在提供有关于订货取消预订的信息之前，你必须从用户处获取以下信息：预定号，客户姓名。
                请讲中文。
                今天的日期是{{current_date}}。
                """)
        TokenStream stream(@UserMessage String message,@V("current_date") String currentDate);
    }

    @Bean
    public Assistant assistant(ChatLanguageModel model, StreamingChatLanguageModel streamModel, ToolsService toolsService, EmbeddingStore embeddingStore, QwenEmbeddingModel qwenEmbeddingModel){
        //对话记忆
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);
        //内容检索器
        EmbeddingStoreContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(qwenEmbeddingModel)
                .maxResults(5)
                .minScore(0.6)
                .build();
        return AiServices.builder(Assistant.class)
                .chatLanguageModel(model)
                .streamingChatLanguageModel(streamModel)
                .chatMemory(chatMemory)
                .tools(toolsService)
                .contentRetriever(contentRetriever)
                .build();
    }

    public interface UniqueAssistant{
        String chat(@MemoryId int memoryId,@UserMessage String message);
        //流式响应
        TokenStream stream(@MemoryId int memoryId,@UserMessage String message);
    }

    @Bean
    public EmbeddingStore embeddingStore(){
        return new InMemoryEmbeddingStore();
    }

    @Bean
    public UniqueAssistant uniqueAssistant(ChatLanguageModel model, StreamingChatLanguageModel streamModel){
        PersistentChatMemoryStore store = new PersistentChatMemoryStore();

        ChatMemoryProvider provider = memoryId -> MessageWindowChatMemory.builder().id(memoryId).maxMessages(10).chatMemoryStore(store).build();
        return AiServices.builder(UniqueAssistant.class)
                .chatLanguageModel(model)
                .streamingChatLanguageModel(streamModel)
                .chatMemoryProvider(provider)
                .build();
    }
}
